import matplotlib.pyplot as plt
import numpy as np

from sklearn.metrics import classification_report
from sklearn.preprocessing import LabelBinarizer
from sklearn.model_selection import train_test_split
from tensorflow.keras.optimizers import SGD
from sklearn.datasets import fetch_openml
from tensorflow.keras import backend

from lenet import Lenet


def main():
    print("[info]:开始加载MNIST数据集...")
    raw_data = fetch_openml('mnist_784')
    data = raw_data.data.astype("float") / 255.0
    train_x, test_x, train_y, test_y = train_test_split(data, raw_data.target.astype("int"),
                                                        test_size=0.25, random_state=42)
    train_x = train_x.values
    test_x = test_x.values
    train_x = train_x.reshape(train_x.shape[0], 28, 28, 1)
    test_x = test_x.reshape(test_x.shape[0], 28, 28, 1)
    le = LabelBinarizer()
    train_y = le.fit_transform(train_y)
    test_y = le.transform(test_y)

    # 编译模型
    model = Lenet.build(width=28, height=28, depth=1, classes=10)
    opt = SGD(0.01)
    model.compile(loss="categorical_crossentropy", optimizer=opt,
                  metrics=["acc"])

    # 训练模型
    print("[info]:开始训练模型...")
    record = model.fit(train_x, train_y, validation_data=(test_x, test_y),
                       batch_size=128, epochs=20, verbose=1)

    # 评估网络
    print("[info]:评估网络中...")
    predictions = model.predict(test_x, batch_size=128, verbose=1)
    print(classification_report(test_y.argmax(1),
                                predictions.argmax(1),
                                target_names=[str(i) for i in le.classes_]))
    # 画图
    plt.style.use("ggplot")
    plt.figure()
    plt.plot(np.arange(0, 20), record.history["loss"], label="train_loss")
    plt.plot(np.arange(0, 20), record.history["val_loss"], label="val_loss")
    plt.plot(np.arange(0, 20), record.history["acc"], label="train_acc")
    plt.plot(np.arange(0, 20), record.history["val_acc"], label="val_acc")
    plt.title("Training Loss and Accuracy")
    plt.xlabel("Epoch #")
    plt.ylabel("Loss/Accuracy")
    plt.legend()
    plt.show()


if __name__ == '__main__':
    main()
